import re

import requests
import json

API_KEY = "vRyL4Nfqxp8MN3W5kIgyKEDV"
SECRET_KEY = "EwRWwW7V271mY0wE3H0TS2AVHCiZkk27"
def get_access_token():
    """
    使用 API Key，Secret Key 获取access_token，替换下列示例中的应用API Key、应用Secret Key
    """

    url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={}&client_secret={}"
    final_url = url.format(API_KEY, SECRET_KEY)
    payload = json.dumps("")
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", final_url, headers=headers, data=payload)
    return response.json().get("access_token")


def main(question):
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-tiny-8k?access_token=" + get_access_token()

    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": question
            }
        ],
        "stream": True
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload, stream=True)
    for line in response.iter_lines():
        if line:
            match = re.search(r'"result":\s*"([^"]+)"', line.decode("UTF-8"))

            if match:
                # 提取匹配到的内容（第一个括号内的部分）
                result_content = match.group(1).replace('\\n', '\n')
                print(result_content,end="")  # 输出：这是你想要的内容


if __name__ == '__main__':
    question = ""
    while (True):
        question = input("罗欣:")
        answer = main(question)
        print("李冰:", answer)